/*
 * Copyright (c) China Telecom Cloud Technology Co., Ltd. 2024-2025. All rights reserved.
 */
#ifndef MEMORY_POOL_H
#define MEMORY_POOL_H

#include <cstdint>
#include <vector>
#include <list>
#include <unordered_map>
#include <mutex>
#include <shared_mutex>
#include <atomic>
#include <set>
#include <memory>
#include "Segment.h"
#include "KvcBaseDefine.h"
#include "LruManager.h"

#if defined(USE_GPU)
#include <cuda_runtime.h>
#endif

// TODO(yangxianpku: 2025.08.21):  移除typedef模式
typedef struct {
    int64_t dramHit;  // KVClient命中次数, QueryPrefixCache时统计
    int64_t total;    // 总查询次数, QueryPrefixCache时统计
} cacheStatis;

using OffsetId = uint64_t;

class MemoryPool {
private:
    void*    _memStartAddr;          // 内存池起始地址

    //! 每个MemoryPool拥有固定的_memSegSize，不同的创建不同的MemoryPool
    uint64_t _memSegSize;              // 内存块大小(业务层可见), 字节
    uint64_t _phyMemSegSize;           // 真实内存块大小(物理层, 会按页大小对齐)， 字节
    uint64_t _poolSize;                // 内存池总大小(业务层)
    uint64_t _phyPoolSize;             // 内存池总物理大小(按页 + 大页对齐)
    uint32_t _blkNum;                  // 总块数

    //! 当初设想做一个多节点共享内存(暂时没使用)
    bool                _isShare;      // 是否共享内存
    std::string         _shmFilePath;  // 共享内存文件路径(绝对路径)

    std::shared_mutex   _memPoolLock;  // memPool对象读写锁(读共享，写互斥)
    LRUManager          _lruMap;
    cacheStatis         _statis;       // 统计命中率

    std::vector<Segment> _segList;     // 全局Segments

    std::list<OffsetId>  _freeList;    // 空闲块列表(uint64_t为内存地址的offset)


    /**
     *! @brief 
     * 
     * @param allocIdxs 
     * @param num 
     * @return int32_t 
     */
    // TODO(yangxianpku: 2025.08.18): 方法未实现
    int32_t allocFromLruList(std::vector<int32_t> &allocIdxs, 
                            int32_t num);


    /**
     *! @brief 从全局segments列表中分配预期数量的segments(必要时发生驱逐)
     * 
     * @param blockHashids      待分配的segments的哈希信息
     * @param policy            分配时使用的驱逐策略
     * @param segments          成功分配后返回的segments
     * @param proposeNum        期望分配的块
     * @param writeInRecordSize 
     * @param cacheInfo         KV Cache块信息
     * @return int32_t 
     */
    // TODO(yangxianpku: 2025.08.18): 分配Segment的初始writeInRecordSize不应该为0么，为什么要传一个？
    int32_t allocateSegmentByList(std::vector<BlockHashId> &blockHashids, 
                                AllocPolicy policy,
                                std::vector<Segment *> &segments, 
                                int32_t proposeNum, 
                                int32_t writeInRecordSize,
                                CacheInfo cacheInfo);
    
    /**
     *! @brief 通过Segment的内存偏移获取其idx
     * 
     * @param offset Segment的内存地址相对内存池的偏移
     * @return int32_t 
     */
    int32_t getIdxByOffset(uint64_t offset);

public:
    MemoryPool() = default;

    /**
     *! @brief 内存池初始化，赋值变量，重置统计信息
     * 
     * @param memStartAddr   内存池起始地址
     * @param memSegSize     内存块大小(业务层可见), 字节
     * @param phyMemSegSize  真实内存块大小(物理层, 会按页大小对齐)， 字节
     * @param poolSize       内存池总大小(业务层)
     * @param phyPoolSize    内存池总物理大小(按页 + 大页对齐)
     * @param blkNum         总块数
     * @param isShare        是否共享内存
     * @param freeList       空闲块列表(uint64_t为内存地址的offset)
     * @param shmFilePath    共享内存文件路径(绝对路径)
     * @param segList        全局Segments
     */
    void init(void *memStartAddr, 
            uint32_t memSegSize, 
            uint32_t phyMemSegSize, 
            uint64_t poolSize, 
            uint64_t phyPoolSize,
            int32_t blkNum, 
            bool isShare, 
            std::list<uint64_t> &freeList, 
            std::string shmFilePath,
            std::vector<Segment> &segList
        );

    /**
     *! @brief 返回共享内存文件路径
     * 
     * @return std::string 
     */
    std::string getShmFilePath();


    /**
     *! @brief 返回内存池物理真实大小
     * 
     * @return uint64_t 
     */
    uint64_t getPhyPoolSize();


    /**
     *! @brief 获取内存块大小(逻辑大小，字节)
     * 
     * @return uint32_t 
     */
    uint32_t getMemSegSize();


    /**
     *! @brief 返回内存池的起始地址
     * 
     * @return void* 地址指针
     */
    void *getMemStartAddr();


    /**
     *! @brief 返回全局Segment列表
     * 
     * @return std::vector<Segment>& 
     */
    std::vector<Segment> &getSeglist();


    /**
     *! @brief 返回空闲Segment列表(返回Segment的offeset非真实的Segment)
     * 
     * @return std::list<uint64_t>&   offset列表
     */
    std::list<uint64_t> &getFreelist();


    /**
     *! @brief 根据segmentID查询Segment
     * 
     * @param segmentId   Segmengt Id
     * @param pSegmentPtr 指向对应的segment的指针的地址
     * @return int32_t    执行成功返回0，失败返回非0
     */
    int32_t getSegmentByID(int64_t segmentId, Segment **pSegmentPtr);


    /**
     *! @brief 判断segmentId对应的Segment是否在内存池中
     * 
     * @param segmentId 
     * @return bool
     */
    bool isInMemoryPool(int64_t segmentId);



    /**
     *! @brief 批量分配Sgements
     * 
     * @param blockHashids 
     * @param policy            分配期间使用的淘汰策略
     * @param segments 
     * @param writeInRecordSize  TP场景下的写入次数(是否就等于TP大小？)
     * @param cacheInfo 
     * @return int32_t 
     */
    int32_t allocateSegments(std::vector<BlockHashId> &blockHashids, 
                            AllocPolicy policy,
                            std::vector<Segment *> &segments, 
                            int32_t writeInRecordSize, 
                            CacheInfo cacheInfo
                        );

    /**
     *! @brief 将segmentId对应的Segment设置为只读
     * 
     * @param segmentId 
     * 
     * @return int32_t 程序运行错误码
     */
    int32_t sealSegment(int64_t segmentId);




    /**
     *! @brief 复用旧Segment(旧Segment的内存使用权转移到新的)
     * 
     * @param oldSegmentId   内存资源提供方
     * @param newSegmentId   内存资源接收方
     * 
     * @return int32_t 执行成功返回具体offset，失败返回全F
     */
    // TODO(yangxianpku: 2025.08.18): 方法未实现
    int32_t reuseSegment(int64_t oldSegmentId, int64_t newSegmentId);



    /**
     *! @brief 释放特定的Segment将其放回空闲列表中
     * 
     * @param segmentId 
     */
    void releaseSegment(int64_t segmentId);


    /**
     *! @brief 按序在MemPool中检索Segment，如果有没找到的，则返回已找到的Segment列表
     * 
     * @param segmentIds 
     * 
     * @return std::vector<Segment>  前N个命中的SegmentId对应的segment
     */
    std::vector<Segment> readInOrder(std::vector<int64_t> &segmentIds);



    /**
     *! @brief 获取所有FilledSegments
     * 
     * @return std::vector<BlockHashId> 
     */
    std::vector<BlockHashId> getAllCachedSegmentIds();


    /*
     * 唯一确定segment的信息
     */
    // TODO(yangxianpku: 2025.08.18): 类型未使用
    typedef struct tagSegmentIdentifier {
        int64_t hashId;
        int32_t tpNum;

        // 重载小于运算符
        bool operator<(const tagSegmentIdentifier &other) const
        {
            if (hashId != other.hashId) {
                return hashId < other.hashId;
            }
            return tpNum < other.tpNum;
        }
    } SegmentIdentifier;


    /**
     *! @brief 将对应的Segment添加到末尾，如果已经存在则移动到末尾
     * 
     * @param segmentIds 
     */
    void reorder(std::vector<int64_t> &segmentIds);

    void printLruStatus();

    void resetCacheStatis();

    void updateCacheStatis(int32_t dramHitDelta, 
                        int32_t totalDelta);

    void printCacheStatis();

    void cleanMemPool();
};


/**
 *! @brief 创建内存池对象
 * 
 * @param poolSize    内存池总大小
 * @param memSegSize  内存块大小
 * @param isShare     是否为共享内存池
 * @param needLock    是否需要锁页物理内存(Pin??)
 * @param pool        指向新创建的内存池的指针
 * 
 * @return int32_t    创建成功返回0, 创建失败返回具体错误码
 */
int32_t createMemoryPool(uint64_t poolSize, 
                        uint32_t memSegSize, 
                        bool isShare, 
                        bool needLock, 
                        MemoryPool *pool
                    );


/**
 *! @brief 销毁内存池释放空间
 * 
 * @param pool 
 */
void destroyMemoryPool(MemoryPool *pool);

#endif  // MEMORY_POOL_H